import pandas as pd
import matplotlib.pyplot as plt
import glob
import os
import matplotlib

# Set the font to Arial
matplotlib.rcParams['font.family'] = 'Arial'

# Directory containing your CSV files
csv_directory = r'C:\Users\jakea\Desktop\new CaM alphafold\cyanobacteria models\csv'

# Find all CSV files in the directory
csv_files = glob.glob(os.path.join(csv_directory, '*.csv'))

# Iterate over each CSV file
for csv_file in csv_files:
    # Extract the file name without directory and extension for use in plot titles and file names
    file_name = os.path.basename(csv_file).replace('.csv', '')

    # Read the CSV file into a DataFrame
    df = pd.read_csv(csv_file, sep=',')

    # Display the first few rows and columns to understand the structure
    print(f"Processing file: {file_name}")
    print(df.head())
    print(df.columns)

    # Adjust the column names based on inspection
    residue_column = 'Residue Number'  # Replace with the actual column name
    plddt_column = 'pLDDT'    # Replace with the actual column name

    # Create a DataFrame with relevant columns
    data = pd.DataFrame({
        'Residue Number': df[residue_column],
        'pLDDT': df[plddt_column]
    })

    # Group by 'Residue Number' and calculate the average pLDDT Score
    average_scores = data.groupby('Residue Number').mean().reset_index()

    # Plot the averaged data
    plt.figure(figsize=(6, 4))
    plt.plot(average_scores['Residue Number'], average_scores['pLDDT'], linestyle='-', linewidth=2, color='black')

    # Set labels and title with specific font sizes
    plt.xlabel('Residue Number', fontsize=16, labelpad=18)
    plt.ylabel('Average pLDDT Score', fontsize=16, labelpad=18)
    plt.title(f'Average AlphaFold pLDDT Scores per Residue - {file_name}', fontsize=16, pad=20)

    # Customize tick parameters for x and y axes
    plt.tick_params(axis='both', which='major', labelsize=14)
    
    # Set y-axis limits
    plt.ylim(0, 100)

    # Add background color spans
    plt.axhspan(0, 50, facecolor='red', alpha=0.3)
    plt.axhspan(50, 70, facecolor='orange', alpha=0.3)
    plt.axhspan(70, 90, facecolor='lightblue', alpha=0.3)
    plt.axhspan(90, 100, facecolor='darkblue', alpha=0.3)

    plt.grid(True)
    #plt.xticks(rotation=90)  # Rotate x-axis labels if needed
    plt.tight_layout()  # Adjust layout to make room for rotated labels

    # Save the plot as an image file
    plt.savefig(os.path.join(csv_directory, f'{file_name}_shaded_pLDDT.png'), dpi=600)
    plt.close()  # Close the plot to avoid memory issues

print("Processing and plotting complete.")
